// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief	'C' Implementation of the PMODEL (BT) Algorithm
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "pmodel_assoc.h"
#include "copyfile.h"
#include "pmaths.h"
#include "utils.h"

#ifndef JUNK_CHISQ
#define JUNK_CHISQ (double) -9.0
#endif

void pmodel_write_file_header(FILE *out, struct selected_options *ops) { 
	if (out != NULL && ops != NULL) {
		int ci = (int) (100 * ops->ci_level);
		if (!ops->model_perm_gen) fprintf(out, "CHR,BP,SNP,TEST,AFF,UNAFF,OR,L%i,U%i,CHISQ,DF,P,EMP1,EMP2,BONF,HOLM,FDR_BH\n", ci, ci); 
		else fprintf(out, "CHR,BP,SNP,TEST,AFF,UNAFF,CHISQ,DF,P,EMP1,EMP2,BONF,HOLM,FDR_BH\n", ci, ci); 
	}
}

/*!
\ingroup pmodel_alg 
\brief Calculation code of PMODEL (BT).
\details
<p>Code based on the PLINK <em>model</em> %option.</p>
<p>This function runs in 2 distinct modes.<br>
It can run in a simple calculation mode or in a result persistence mode
where data is written to file.<br>
To switch on the persistence mode, the FILE output handle and the adjusted parameters
must not be NULL.</p>

<p>When in calculation mode, the FILE output handle, adjusted struct, EMP1/EMP2 parameters can be set to NULL 
or zero.</p>

<p>%Model calculations performed are specified by setting the appropriate flag bound 
to the options structure (\ref selected_options) to TRUE.</p>

<p>Following the logic of PLINK, if only a single model %option is specified then this functions returns a CHISQ value for that calculation.<br>
If more than 1 model %option is specified, this function returns 1 - Best Pvalue for all the calculations run.
</p>

\param [in] X Flag that SNP on sex chromsome
\param [in] haploid Flag that SNP is haploid
\param [in] ones_and_twos BitString containing the allelic data of a SNP
\param [in] samples Data-set %sample array (people and phenotypes)
\param [in] nSamples Size of the %sample array
\param [in] ops Models to run.
\param [in] snp SNP under investigation
\param [in,out] out Output file handle
\param [in] emp1 Empirical P value ratio for a SNP
\param [in] emp2 %Family Wise Error Rate (FWER) for a SNP
\param [in] adj Structure containing the adjusted p-values (BONF, HOLM....)
\param [in] zt Normalised confidence interval threshold (zt = ltqnorm(1 - (1 - ops.ci_level) / 2))
\returns
	double or JUNK_CHISQ is calculation went mental.
\sa Plink::fullModelAssoc()
\sa selected_options::model_perm_gen 
\sa selected_options::model_perm_dom 
\sa selected_options::model_perm_rec 
\sa selected_options::model_perm_trend 
\sa selected_options::model_perm_allelic
*/
double pmodel_assoc(BOOL X, BOOL haploid, BitString ones_and_twos, struct sample *samples, int nSamples, struct selected_options *ops, struct small_locus *snp, FILE *out, double emp1, double emp2, struct adjusted *adj, double zt) {
	int A11 = 0, A12 = 0, A22 = 0, U11 = 0, U12 = 0, U22 = 0;
	register int i, modelCount = 0;
	double best_p = 1;
	
	// Association statistics
	double obs_A, obs_U, obs_T;
	double obs_1, obs_2;
	double obs_11, obs_12, obs_22;
	BOOL invalid = FALSE;

	// Cochram-Armitage Trend test variables.
	double CA = 0, varCA = 0, CA_chisq = 0, CA_p = 0.0;

	// Multiplicative model variables.
	double mult_p, mult_chisq;
	double obs_A1, obs_A2;
	double obs_U1, obs_U2;
	double exp_A1, exp_A2, exp_U1, exp_U2;

	// Result variables.
	double gen_p, dom_p, rec_p;
	double dom_chisq = 0, rec_chisq = 0, gen_chisq = 0;

	if (haploid) return 0; // Skip haploid markers
	
	// Models number (effects the return value).
	if (ops->model_perm_allelic) modelCount++;
	if (ops->model_perm_dom) modelCount++;
	if (ops->model_perm_gen) modelCount++;
	if (ops->model_perm_rec) modelCount++;
	if (ops->model_perm_trend) modelCount++;

	// Iterate over individuals
	for (i = 0; i < nSamples; i++) {
		unsigned int mask;
		long bitpos = (long) i * 2;
		BOOL s1 = FALSE, s2 = FALSE;

		if (samples[i].pperson->missing) continue;

		// Read S1 value
		mask = 0x1;
		mask = mask << (bitpos % BITBULKSIZE);
		mask = ones_and_twos[bitpos/BITBULKSIZE] & mask;
		if (mask !=0) s1 = TRUE;
		bitpos++;

		// Read S2 value
		mask = 0x1;
		mask = mask << (bitpos % BITBULKSIZE);
		mask = ones_and_twos[bitpos/BITBULKSIZE] & mask;
		if (mask !=0) s2 = TRUE;

		// Only consider diploid chromosomes
		if (!(X && samples[i].sex)) {
			if (samples[i].pperson->aff) { // cases
				if (!s1) {
					if (!s2) A11++; // Homozyg 00
					else A12++;     // Hetero  01            
				}
				else if (s2) A22++; // Homozyg 11
			} 
			else {
				if (!s1) { 
					if (!s2) U11++;   // Homozyg 00
					else U12++;       // Hetero  01
				} else if (s2) U22++; // Homozyg 11
			}
		}
	}

	// Calculate association statistics
	obs_A = (double) (A11 + A12 + A22);
	obs_U = (double) (U11 + U12 + U22);
	obs_T = (double) (obs_A + obs_U);

	obs_1 = (double) ((2 * (A11+U11)) + A12 + U12);
	obs_2 = (double) ((2 * (A22+U22)) + A12 + U12);

	obs_11 = (double) (A11 + U11);
	obs_12 = (double) (A12 + U12);
	obs_22 = (double) (A22+ U22);

	if (A11 < ops->min_geno_cell || A12 < ops->min_geno_cell || A22 < ops->min_geno_cell) 
		invalid = TRUE;
	else if (U11 < ops->min_geno_cell || U12 < ops->min_geno_cell || U22 < ops->min_geno_cell) 
		invalid = TRUE;

	// Cochram-Armitage Trend test
	if (ops->model_perm_trend) {
		CA = ( ( obs_U / obs_T * A12 ) - ( obs_A / obs_T * U12 ) ) + 2*( ( obs_U / obs_T * A22 ) - ( obs_A / obs_T * U22 ) );
		varCA = obs_A * obs_U * ( ( obs_T * ( obs_12 + 4*obs_22 ) - ( obs_12+2*obs_22 ) * ( obs_12+2*obs_22 )  )  / (obs_T * obs_T * obs_T  )) ;
		CA_chisq = (CA*CA) / varCA;
		CA_p = chiprobP(CA_chisq,1);
	}

	// Multiplicative model
	obs_A1 = (double) (2*A11 + A12);
	obs_A2 = (double) (2*A22 + A12);
	obs_U1 = (double) (2*U11 + U12);
	obs_U2 = (double) (2*U22 + U12);

	// note 2's cancelled for obs_A and obs_T 
	// which are counts of individuals, not alleles.
	exp_A1 = (obs_A * obs_1 ) / obs_T;
	exp_A2 = (obs_A * obs_2 ) / obs_T;
	exp_U1 = (obs_U * obs_1 ) / obs_T;
	exp_U2 = (obs_U * obs_2 ) / obs_T;

	mult_chisq =  ( ( obs_A1 - exp_A1 ) * ( obs_A1 - exp_A1 ) ) / exp_A1
		+ ( ( obs_A2 - exp_A2 ) * ( obs_A2 - exp_A2 ) ) / exp_A2
		+ ( ( obs_U1 - exp_U1 ) * ( obs_U1 - exp_U1 ) ) / exp_U1
		+ ( ( obs_U2 - exp_U2 ) * ( obs_U2 - exp_U2 ) ) / exp_U2;

	mult_p = chiprobP(mult_chisq,1);
	gen_p = dom_p = rec_p = -9;

	if (!invalid) {
		// General model (Chi-Test)
		double exp_A11 = (obs_A * obs_11 ) / obs_T;
		double exp_A12 = (obs_A * obs_12 ) / obs_T;
		double exp_A22 = (obs_A * obs_22 ) / obs_T;
		double exp_U11 = (obs_U * obs_11 ) / obs_T;
		double exp_U12 = (obs_U * obs_12 ) / obs_T;
		double exp_U22 = (obs_U * obs_22 ) / obs_T;

		// Genotypic model
		if (ops->model_perm_gen) {
			gen_chisq =  ( ( A11 - exp_A11 ) * ( A11 - exp_A11 ) ) / exp_A11
				+ ( ( A12 - exp_A12 ) * ( A12 - exp_A12 ) ) / exp_A12
				+ ( ( A22 - exp_A22 ) * ( A22 - exp_A22 ) ) / exp_A22
				+ ( ( U11 - exp_U11 ) * ( U11 - exp_U11 ) ) / exp_U11
				+ ( ( U12 - exp_U12 ) * ( U12 - exp_U12 ) ) / exp_U12
				+ ( ( U22 - exp_U22 ) * ( U22 - exp_U22 ) ) / exp_U22;
			gen_p = chiprobP(gen_chisq,2);
		}

		// Dominant (minor allele) (1) model
		if (ops->model_perm_dom) {
			dom_chisq =  ( ( (A11+A12) - (exp_A11+exp_A12) ) * ( (A11+A12) - (exp_A11+exp_A12) ) ) / (exp_A11+exp_A12) 
				+ ( ( A22 - exp_A22 ) * ( A22 - exp_A22 ) ) / exp_A22
				+ ( ( (U11+U12) - (exp_U11+exp_U12) ) * ( (U11+U12) - (exp_U11+exp_U12) ) ) / (exp_U11+exp_U12) 
				+ ( ( U22 - exp_U22 ) * ( U22 - exp_U22 ) ) / exp_U22;
			dom_p = chiprobP(dom_chisq,1);
		}

		// Recessive (minor allele) (1) model
		if (ops->model_perm_rec) {
			rec_chisq =  ( ( (A22+A12) - (exp_A22+exp_A12) ) * ( (A22+A12) - (exp_A22+exp_A12) ) ) / (exp_A22+exp_A12) 
				+ ( ( A11 - exp_A11 ) * ( A11 - exp_A11 ) ) / exp_A11
				+ ( ( (U22+U12) - (exp_U22+exp_U12) ) * ( (U22+U12) - (exp_U22+exp_U12) ) ) / (exp_U22+exp_U12) 
				+ ( ( U11 - exp_U11 ) * ( U11 - exp_U11 ) ) / exp_U11;
			rec_p = chiprobP(rec_chisq,1);
		}
	}

	// Save the desired result
	best_p = mult_p;
	if (!invalid) {
		if (gen_p < best_p && gen_p >= 0 && ops->model_perm_gen) best_p = gen_p; // general
		if (dom_p < best_p && dom_p >= 0 && ops->model_perm_dom) best_p = dom_p; // dom
		if (rec_p < best_p && rec_p >= 0 && ops->model_perm_rec) best_p = rec_p; // rec
	}

	// Write initial model results to output file.
	if (out != NULL && adj != NULL) {
		int v1, v2, v3, v4;
		double OR = 0.0, lowerCI = 0.0, upperCI = 0.0;

		// Genotypic
		if (ops->model_perm_gen) {
			fprintf(out, "%i,%i,%s,%s,%i/%i/%i,%i/%i/%i,", snp->chr, snp->bp, snp->name, "GENO", A11, A12, A22, U11, U12, U22);
			if (gen_p < -1) fprintf(out, "NA,NA,NA,1,1,");
			else fprintf(out, "%f,%s,%f,%f,%f,", gen_chisq, "2", gen_p, emp1, emp2);
		}

		// CA Trend test
		if (ops->model_perm_trend) {
			v1 = (A11 * 2) + A12;
			v2 = A12 + (A22 * 2);
			v3 = (U11 * 2) + U12;
			v4 = U12 + (U22 * 2);
			getOddsRatio(v1, v2, v3, v4, &OR, &lowerCI, &upperCI, zt);
			fprintf(out, "%i,%i,%s,%s,%i/%i,%i/%i,%f,%f,%f,", snp->chr, snp->bp, snp->name, "TREND", v1, v2, v3, v4, OR, lowerCI, upperCI);
			if (CA_p < -1) fprintf(out, "NA,NA,NA,1,1,");
			else fprintf(out, "%f,%s,%f,%f,%f,", CA_chisq, "1", CA_p, emp1, emp2);
		}

		// Allelic
		if (ops->model_perm_allelic) {
			v1 = (A11 * 2) + A12;
			v2 = A12 + (A22 * 2);
			v3 = (U11 * 2)+ U12;
			v4 = U12+ (U22 * 2);
			getOddsRatio(v1, v2, v3, v4, &OR, &lowerCI, &upperCI, zt);
			fprintf(out, "%i,%i,%s,%s,%i/%i,%i/%i,%f,%f,%f,", snp->chr, snp->bp, snp->name, "ALLELIC", v1, v2, v3, v4, OR, lowerCI, upperCI);
			if (mult_p < -1) fprintf(out, "NA,NA,NA,1,1,");
			else fprintf(out, "%f,%s,%f,%f,%f,", mult_chisq, "1", mult_p, emp1, emp2);
		}

		// Dominant
		if (ops->model_perm_dom) {
			v1 = A11 + A12;
			v2 = A22;
			v3 = U11 + U12;
			v4 = U22;
			getOddsRatio(v1, v2, v3, v4, &OR, &lowerCI, &upperCI, zt);
			fprintf(out, "%i,%i,%s,%s,%i/%i,%i/%i,%f,%f,%f,", snp->chr, snp->bp, snp->name, "DOM", v1, v2, v3, v4, OR, lowerCI, upperCI);
			if (dom_p < -1) fprintf(out, "NA,NA,NA,1,1,");
			else fprintf(out, "%f,%s,%f,%f,%f,", dom_chisq, "1", dom_p, emp1, emp2);
		}

		// Recessive
		if (ops->model_perm_rec) {
			v1 = A11;
			v2 = A12 + A22;
			v3 = U11;
			v4 = U12 + U22;
			getOddsRatio(v1, v2, v3, v4, &OR, &lowerCI, &upperCI, zt);
			fprintf(out, "%i,%i,%s,%s,%i/%i,%i/%i,%f,%f,%f,", snp->chr, snp->bp, snp->name, "REC", v1, v2, v3, v4, OR, lowerCI, upperCI);
			if (rec_p < -1) fprintf(out, "NA,NA,NA,1,1,");
			else fprintf(out, "%f,%s,%f,%f,%f,", rec_chisq, "1", rec_p, emp1, emp2);
		} 

		fprintf(out, "%f,%f,%f\n", adj->bonf, adj->holm, adj->fdr_bh);
	}

	if (modelCount == 1) {
		double chisq = 0;

		if (ops->model_perm_gen) chisq = gen_p >= 0 ? gen_chisq : JUNK_CHISQ;
		else if (ops->model_perm_dom) chisq = dom_p >= 0 ? dom_chisq : JUNK_CHISQ;
		else if (ops->model_perm_rec) chisq = rec_p >= 0 ? rec_chisq : JUNK_CHISQ;
		else if (ops->model_perm_trend) chisq = CA_p >= 0 ? CA_chisq : JUNK_CHISQ;
		else if (ops->model_perm_allelic) chisq = mult_p >= 0 ? mult_chisq : JUNK_CHISQ; // Basic allelic test

		return chisq;
	} else 
		return 1 - best_p;
}

#define SUFFIX_MODEL_OPTIONS ".options.txt"

int pmodel_write_model_options(struct selected_options *ops) {
	FILE *out = NULL;
	BOOL jobDone = FALSE;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];

	if (ops == NULL) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0 || strlen(ops->rdir) == 0) return 0;

	sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_MODEL_OPTIONS);
	sprintf(remoteFilepath, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_MODEL_OPTIONS);
	
	out = fopen(localFilepath, "w");
	if (out != NULL) {
		fprintf(out, "%i,%i,%i,%i,%i\n", ops->model_perm_dom, ops->model_perm_rec, ops->model_perm_trend, ops->model_perm_gen, ops->model_perm_allelic);
		fclose(out);
		if (copyAndDeleteFile(localFilepath, remoteFilepath)) jobDone = TRUE;
	}

	if (jobDone) return 1;
	else return 0;
}

int pmodel_read_model_options(struct selected_options *ops) {
	FILE *in = NULL;
	BOOL jobDone = FALSE;
	char remoteFilepath[FILE_PATH_LENGTH];
	
	if (ops == NULL) return 0;
	else if (strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0 || strlen(ops->rdir) == 0) return 0;

	sprintf(remoteFilepath, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_MODEL_OPTIONS);
	in = fopen(remoteFilepath, "r");
	if (in != NULL) {
		BOOL readFile = FALSE;
		char line[1024];
		if (fgets(line, 1024, in)) readFile = TRUE;
		fclose(in);
		if (readFile) {
			if (strlen(line) > 0) {
				char token[80];

				get_token(line, token, 0, ',');
				ops->model_perm_dom = (BOOL) atoi(token);

				get_token(line, token, 1, ',');
				ops->model_perm_rec = (BOOL) atoi(token);

				get_token(line, token, 2, ',');
				ops->model_perm_trend = (BOOL) atoi(token);

				get_token(line, token, 3, ',');
				ops->model_perm_gen = (BOOL) atoi(token);

				get_token(line, token, 4, ',');
				ops->model_perm_allelic = (BOOL) atoi(token);

				jobDone = TRUE;
			}
		}
	}

	if (jobDone) return 1;
	else return 0;
}